from llama_index.llms.openai import OpenAI
import time
import os
from llama_index.core.llms import ChatMessage, TextBlock, ImageBlock
from llama_index.core.tools import FunctionTool


def generate_poem(poet: str, theme: str):
    """根据诗人和主题生成对应的诗歌"""
    print("poet", poet)
    print("theme", theme)
    return {
        "poet": poet,
        "theme": theme,
        "message": f"作者{poet}为{theme}主题创作的诗歌",
    }


# 把generate_poem包装成工具
tool = FunctionTool.from_defaults(fn=generate_poem)
llm = OpenAI(model="gpt-4o")
# predict_and_call根据用户请求选择并调用合适的工具
response = llm.predict_and_call([tool], "以李白口吻编写一首赞美春天的诗歌")
print(str(response))
